def bce_loss(y_true, y_pred_logits):
    # y_true：真实标签（0 或 1）。
    # y_pred_logits：模型的原始预测值（logits），即未经过激活函数的输出。
    # y_pred_prob = F.sigmoid(y_pred_logits)
    # y_true_f = y_true.view(-1)
    # y_true_f = y_true
    y_pred_logits = y_pred_logits.view(-1)
    y_pred_prob_f = y_pred_logits.clamp(min=1e-7, max=1 - 1e-7)  # 调整预测值范围：将预测值限制在 [1e-7, 1 - 1e-7] 之间，以防止对数计算时出现数值不稳定。
    return -(y_true * y_pred_prob_f.log() + (1. - y_true) * (1 - y_pred_prob_f).log()).mean()  # 计算交叉熵损失，并返回平均损失。

class Loss_target(nn.Module):
    def __init__(self, weight_angle=10):
        # 角度损失的权重系数。默认为 10，用于控制角度损失在总损失中的影响程度。
        super(Loss_target, self).__init__()
        self.weight_angle = weight_angle
        self.bce = bce_loss
        # 指定使用的二分类交叉熵损失函数 (bce_loss)，它在 forward 方法中用来计算领域损失。
 
    def forward(self, gt_score, pred_score, gt_geo, pred_geo, ignored_map,pre_domain):
        if torch.sum(gt_score) < 1:
            return torch.sum(pred_score + pred_geo) * 0
        # 如果 gt_score 的总和小于 1（表示没有有效的目标），则返回 0 作为损失，避免计算损失时出现错误或不必要的计算。
 
        classify_loss = get_dice_loss(gt_score, pred_score*(1-ignored_map))
        # 使用 get_dice_loss 计算分类损失。pred_score * (1 - ignored_map) 是预测分数与忽略区域掩码结合后的结果，确保在计算损失时忽略掉指定的区域。
 
        gt_doamin = torch.Tensor([[[1.]]]).to(torch.device("cuda"))
        doamin_loss = self.bce(gt_doamin, pre_domain)
        # 生成一个目标领域分数 gt_doamin，设为 1，并将其转移到 GPU 设备上。使用 bce_loss 计算领域损失，这里领域损失衡量模型预测的领域分数与目标领域分数之间的差异。
 
        iou_loss_map, angle_loss_map = get_geo_loss(gt_geo, pred_geo)
 
        angle_loss = torch.sum(angle_loss_map * gt_score) / torch.sum(gt_score)
        iou_loss = torch.sum(iou_loss_map * gt_score) / torch.sum(gt_score)
        geo_loss = self.weight_angle * angle_loss + iou_loss
        # 使用 get_geo_loss 计算几何损失，包括交并比（IoU）损失和角度损失。
        # 计算角度损失和 IoU 损失，并根据 weight_angle 进行加权，得到总的几何损失。
        # print('classify loss is {:.8f}, angle loss is {:.8f}, iou loss is {:.8f}'.format(classify_loss, angle_loss, iou_loss))
        return geo_loss, classify_loss, doamin_loss  # 返回三个损失值：几何损失、分类损失和领域损失。

cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']  # 配置列表，指定网络层的类型和数量。
 
def make_layers(cfg, batch_norm=False):
    layers = []  # 列表用于存储创建的层。
    in_channels = 3  # 初始化为 3，对应输入图像的通道数。
    for v in cfg:
        if v == 'M':  # 如果值为 'M'，添加一个 2x2 的最大池化层。
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:  # 否则，创建一个卷积层（nn.Conv2d），后跟可选的批归一化层和 ReLU 激活函数。
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v  # 更新 in_channels 为当前卷积层的输出通道数 v。
    return nn.Sequential(*layers)  # 使用 nn.Sequential 将所有层打包成一个顺序容器，返回构造好的网络层。
 
 
class VGG(nn.Module):
    def __init__(self, features):
        super(VGG, self).__init__()
        self.features = features  # self.features：存储卷积层。
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))  # 自适应平均池化层，将特征图大小调整为 7x7。
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 1000),
        )  # 全连接层的序列，包括两个 ReLU 激活层和 Dropout 层，用于分类。
 
        for m in self.modules():  # 初始化权重
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)
 
    def forward(self, x):
        x = self.features(x)  # 通过 self.features 处理输入 x。
        x = self.avgpool(x)  # 应用自适应平均池化，将特征图调整为 7x7。
        x = x.view(x.size(0), -1)  # 将池化后的特征图展平，以适应全连接层输入。
        x = self.classifier(x)  # 通过 self.classifier 进行最终的分类预测。
        return x
 
 
class extractor(nn.Module):  # 特征提取层。
    def __init__(self, pretrained):
        super(extractor, self).__init__()
        vgg16_bn = VGG(make_layers(cfg, batch_norm=True))  # 使用 VGG 类创建带有批量归一化的 VGG16 网络。
        if pretrained:
            # state_dict = load_state_dict_from_url('https://download.pytorch.org/models/vgg16-397923af.pth')
            # vgg16_bn.load_state_dict(state_dict)
            if pretrained:  # 如果 pretrained 为 True，则加载本地路径的预训练权重。
                vgg16_bn.load_state_dict(torch.load(
                    '/data/data_weijiawu/Sence_Text_detection/Paper-ACCV/baseline/EAST_1/vgg16_bn-6c64b313.pth'))
        self.features = vgg16_bn.features  # self.features 被设置为 VGG16 网络的卷积层部分。(使用 VGG16 进行特征提取)
 
    def forward(self, x):
        out = []
        for m in self.features:  # 遍历 self.features 中的每一层，对输入 x 进行逐层处理。
            x = m(x)
            if isinstance(m, nn.MaxPool2d):  # 将每个最大池化层的输出保存到 out 列表中。
                out.append(x)
        return out[1:]  # 返回 out[1:]，即从第二个池化层开始的特征图。这通常是为了匹配特定的特征提取需求。

class merge(nn.Module):  # EAST 的核心特征融合层。负责在特征图上进行一系列卷积操作并融合来自不同尺度的特征。它的主要任务是通过逐步上采样和融合不同层的特征图来生成最终的输出特征图。这通常用于目标检测和分割任务中的特征融合阶段。
    def __init__(self):
        super(merge, self).__init__()
 
        self.conv1 = nn.Conv2d(1024, 128, 1)
        self.bn1 = nn.BatchNorm2d(128)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(128, 128, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(128)
        self.relu2 = nn.ReLU()
 
        self.conv3 = nn.Conv2d(384, 64, 1)
        self.bn3 = nn.BatchNorm2d(64)
        self.relu3 = nn.ReLU()
        self.conv4 = nn.Conv2d(64, 64, 3, padding=1)
        self.bn4 = nn.BatchNorm2d(64)
        self.relu4 = nn.ReLU()
 
        self.conv5 = nn.Conv2d(192, 32, 1)
        self.bn5 = nn.BatchNorm2d(32)
        self.relu5 = nn.ReLU()
        self.conv6 = nn.Conv2d(32, 32, 3, padding=1)
        self.bn6 = nn.BatchNorm2d(32)
        self.relu6 = nn.ReLU()
 
        self.conv7 = nn.Conv2d(32, 32, 3, padding=1)
        self.bn7 = nn.BatchNorm2d(32)
        self.relu7 = nn.ReLU()
 
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
 
    def forward(self, x):
        # F.interpolate：将特征图上采样到较大的空间分辨率，使用 bilinear 插值方法。
        # torch.cat：将上采样后的特征图与来自不同层的特征图沿通道维度拼接，结合不同尺度的信息。
        y = F.interpolate(x[3], scale_factor=2, mode='bilinear', align_corners=True)
        y = torch.cat((y, x[2]), 1)
        y = self.relu1(self.bn1(self.conv1(y)))
        y = self.relu2(self.bn2(self.conv2(y)))
 
        y = F.interpolate(y, scale_factor=2, mode='bilinear', align_corners=True)
        y = torch.cat((y, x[1]), 1)
        y = self.relu3(self.bn3(self.conv3(y)))
        y = self.relu4(self.bn4(self.conv4(y)))
 
        y = F.interpolate(y, scale_factor=2, mode='bilinear', align_corners=True)
        y = torch.cat((y, x[0]), 1)
        y = self.relu5(self.bn5(self.conv5(y)))
        y = self.relu6(self.bn6(self.conv6(y)))
 
        y = self.relu7(self.bn7(self.conv7(y)))
        return y
 
 
class output(nn.Module):  # 输出层
    def __init__(self, scope=512):
        super(output, self).__init__()
        self.conv1 = nn.Conv2d(32, 1, 1)  # 卷积层，用于处理输入的特征图。它们的输出通道数不同，目的是生成不同类型的预测。
        self.sigmoid1 = nn.Sigmoid()  # 分别应用于卷积层输出的 Sigmoid 激活函数，用于输出概率。
 
        # self.conv1_1 = nn.Conv2d(32, 1, 1)
        # self.sigmoid1_1 = nn.Sigmoid()
 
        self.conv2 = nn.Conv2d(32, 4, 1)
        self.sigmoid2 = nn.Sigmoid()
        self.conv3 = nn.Conv2d(32, 1, 1)
        self.sigmoid3 = nn.Sigmoid()
        self.scope = 512  # 一个用于缩放位置预测的参数，设为 512。
 
        # GRL domina classifier
        # self.gap = nn.AdaptiveAvgPool2d(1)
        # self.sigmoid_gap = nn.Sigmoid()
        # self.GRL = RevGrad()
        # self.conv_gap = nn.Conv2d(32, 1, 1)
 
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')  # 卷积层的权重使用 Kaiming 正态分布初始化，偏置项初始化为 0。
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
 
    def forward(self, x):
        # score_1 = self.GRL(x)
        # grl = self.sigmoid_gap(self.conv_gap(score_1))
 
        # domain_feature = self.gap(self.GRL(C[3]))
        # class_domain = self.sigmoid_gap(self.conv_gap(domain_feature))
 
        score = self.sigmoid1(self.conv1(x))  # 第一个卷积层的输出，经过 Sigmoid 激活函数处理，代表得分图。
        loc = self.sigmoid2(self.conv2(x)) * self.scope  # 第二个卷积层的输出，经过 Sigmoid 激活函数处理并缩放，表示位置预测。
        angle = (self.sigmoid3(self.conv3(x)) - 0.5) * math.pi  # 第三个卷积层的输出，经过 Sigmoid 激活函数处理并调整，表示角度预测。
        geo = torch.cat((loc, angle), 1)  # 将位置和角度预测拼接起来，形成最终的几何输出。
        return score, geo

class DAImgHead(nn.Module):
    """
    Adds a simple Image-level Domain Classifier head 添加一个图像级领域分类器头，用于图像层面的分类。
    """
 
    def __init__(self):
        """
        Arguments:
            in_channels (int): number of channels of the input feature 输入特征的通道数
            USE_FPN (boolean): whether FPN feature extractor is used 是否使用 FPN 特征提取器
        """
        super(DAImgHead, self).__init__()
 
        self.conv1_da = nn.Conv2d(512, 256, kernel_size=1, stride=1)  # 一个卷积层，将输入通道数从 512 转换为 256。
        self.conv2_da = nn.Conv2d(256, 1, kernel_size=1, stride=1)  # 一个卷积层，将通道数从 256 转换为 1。
 
        for l in [self.conv1_da, self.conv2_da]:  # 权重初始化为均值为 0，标准差为 0.001 的正态分布，偏置初始化为 0。
            torch.nn.init.normal_(l.weight, std=0.001)
            torch.nn.init.constant_(l.bias, 0)
 
    def forward(self, x):  # 输入 x 经过 conv1_da 和 ReLU 激活函数后，再经过 conv2_da 得到输出。
        t = F.relu(self.conv1_da(x))
        return self.conv2_da(t)
 
 
class DAInsHead(nn.Module):
    """
    添加一个实例级领域分类器头，用于实例级别的分类。
    Adds a simple Instance-level Domain Classifier head
    """
 
    def __init__(self):
        """
        Arguments:
            in_channels (int): number of channels of the input feature
        """
        super(DAInsHead, self).__init__()
        self.fc1_da = nn.Linear(2028, 1024)  # 一个全连接层，将 2028 个输入特征转换为 1024 个输出特征。
        self.fc2_da = nn.Linear(1024, 1024)  # 一个全连接层，将 1024 个输入特征转换为 1024 个输出特征。
        self.fc3_da = nn.Linear(1024, 1)  # 一个全连接层，将 1024 个输入特征转换为 1 个输出特征。
        for l in [self.fc1_da, self.fc2_da]:  # 权重初始化为均值为 0，标准差为 0.01 的正态分布（对于 fc1_da 和 fc2_da），fc3_da 的权重标准差为 0.05，偏置初始化为 0。
            nn.init.normal_(l.weight, std=0.01)
            nn.init.constant_(l.bias, 0)
        nn.init.normal_(self.fc3_da.weight, std=0.05)
        nn.init.constant_(self.fc3_da.bias, 0)
 
    def forward(self, x):  # 输入 x 先通过 fc1_da 和 ReLU 激活，然后经过 50% 的 dropout，再通过 fc2_da 和 ReLU 激活，再次经过 50% 的 dropout，最后通过 fc3_da 得到输出。
        x = F.relu(self.fc1_da(x))
        x = F.dropout(x, p=0.5, training=self.training)
 
        x = F.relu(self.fc2_da(x))
        x = F.dropout(x, p=0.5, training=self.training)
 
        x = self.fc3_da(x)
        return x
 
 
class DomainAdaptationModule(torch.nn.Module):
    """
    这个模块的主要作用是通过领域自适应来处理图像特征，使得模型能够更好地适应不同的领域。
    Module for Domain Adaptation Component. Takes feature maps from the backbone and instance
    feature vectors, domain labels and proposals. Works for both FPN and non-FPN.
    域适配组件模块。从主干和实例特征向量、域标签和建议中提取特征图。适用于 FPN 和非 FPN。
    """
 
    def __init__(self):
        super(DomainAdaptationModule, self).__init__()
 
        self.avgpool = nn.AvgPool2d(kernel_size=7, stride=7)  # 一个 7x7 的平均池化层，将输入的特征图缩小到 1x1。
        self.gap = nn.AdaptiveAvgPool2d(1)  # 自适应平均池化层，将输入特征图压缩成每个通道一个值。
        self.sigmoid_gap = nn.Sigmoid()  # 一个 Sigmoid 激活函数，将特征值压缩到 [0, 1] 范围。
 
        self.grl_img = GradientScalarLayer(-1.0 * 0.1)  # 应用梯度反转层，用于在反向传播时反转梯度，鼓励网络学习领域不变特征。
        self.grl_ins = GradientScalarLayer(-1.0 * 0.1)
        self.grl_img_consist = GradientScalarLayer(1.0 * 0.1)  # 与 grl_img 和 grl_ins 类似，但用于不同的目标。
        self.grl_ins_consist = GradientScalarLayer(1.0 * 0.1)
 
        self.imghead = DAImgHead()  # 处理图像级特征的网络头部，用于领域分类。
        self.inshead = DAInsHead()  # 处理实例级特征的网络头部，用于领域分类。
 
 
    def forward(self, img_features):
        """
        Arguments:
            img_features (list[Tensor]): features computed from the images that are used for computing the predictions. 从用于计算预测的图像中计算出的特征。
            da_ins_feature (Tensor): instance-level feature vectors  实例级特征向量。
            da_ins_labels (Tensor): domain labels for instance-level feature vectors  实例级特征向量目标的域标签
            targets (list[BoxList): ground-truth boxes present in the image (optional) 实例级特征向量目标的域标签（list[BoxList）：图像中存在的地面真值框（可选）
 
        Returns:
            losses (dict[Tensor]): the losses for the model during training. During testing, it is an empty dict. 模型在训练过程中的损失。在测试过程中，它是一个空字典。
        """
        img_grl_fea = self.grl_img(self.gap(img_features))  # 经过 grl 的 image-level feature  # 通过 gap 层将特征图压缩为 1x1，然后通过 grl_img 反转梯度。
 
        da_img_features = self.sigmoid_gap(self.imghead(img_grl_fea))  # 对 image-level feature 的每个层都压成 1*1 # 通过 imghead 处理反转后的图像特征，再经过 sigmoid_gap 激活函数，将每个特征值压缩到 [0, 1] 范围。
 
        return da_img_features

class EAST(nn.Module):  # EAST 继承自 nn.Module，这是 PyTorch 中所有神经网络模型的基类。
    def __init__(self, pretrained=True):
        super(EAST, self).__init__()
        self.extractor = extractor(pretrained)  # 初始化一个特征提取器。extractor 可能是一个预训练的卷积神经网络或类似模块，用于从输入图像中提取特征。如果 pretrained=True，则可能会使用预训练权重。
        self.merge = merge()  # 初始化一个合并模块。merge 可能是一个用于处理提取到的特征的模块。
        self.output = output()  # 初始化一个输出模块。output 可能是一个用于生成最终预测的模块。
 
        self.DomainAdaptation = DomainAdaptationModule()  # 初始化一个领域自适应模块。DomainAdaptationModule 可能用于对特征进行领域自适应处理，以适应不同的数据分布或任务。
 
    def forward(self, x):  # forward 方法定义了前向传播的计算过程，即如何将输入数据 x 通过网络进行处理，得到模型的输出。
        C = self.extractor(x)  # 将输入 x 传递给特征提取器，得到特征 C。C 可能是一个包含多个特征图的列表或字典，其中 C[3] 是其中的一个特征图。
        class_domain = self.DomainAdaptation(C[3])  # 将特征图 C[3] 传递给领域自适应模块，得到 class_domain。这是对特征图进行领域自适应处理后的结果。
 
        return self.output(self.merge(C)), class_domain
        # self.merge(C): 将特征 C 传递给合并模块，合并不同层的特征。
        # self.output(self.merge(C)): 将合并后的特征传递给输出模块，生成最终预测。
        # return self.output(self.merge(C)), class_domain：返回两个值，分别是模型的最终预测结果和领域自适应模块的输出。